# !/usr/bin/env python
# -*- coding: utf-8 -*-
"""
__title__ = ''
"""
import torch
from torch import nn
import torch.nn.functional as F
import torch.utils.data as Data
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
import numpy as np
from torchvision import datasets, transforms
from torch.nn import init

# 创建fake data
# torch.manual_seed(99)
# x = torch.unsqueeze(torch.linspace(-1, 1, 1000), dim=1)
# y = x.pow(2) + 0.1 * torch.normal(torch.zeros(x.size()))
# plt.scatter(x.numpy(), y.numpy())
# plt.show()

np.random.seed(666)
X = np.linspace(-1, 1, 1000)
y = np.power(X, 2) + 0.1 * np.random.normal(0, 1, X.size)
print(X.shape)
print(y.shape)
# plt.scatter(X, y)
# plt.show()

# 创建训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, random_state=1024)
X_train = torch.from_numpy(X_train).type(torch.FloatTensor)
X_train = torch.unsqueeze(X_train, dim=1)  # 转换成二维
y_train = torch.from_numpy(y_train).type(torch.FloatTensor)
y_train = torch.unsqueeze(y_train, dim=1)
print(X_train.type)
X_test = torch.from_numpy(X_test).type(torch.FloatTensor)
X_test = torch.unsqueeze(X_test, dim=1)  # 转换成二维

# train_size = int(0.7 * len(X))
# test_size = len(X) - train_size
# X_train, X_test = Data.random_split(X, [train_size, test_size])
# print(len(X_train), len(X_test))


BATCH_SIZE = 50
LR = 0.02
EPOCH = 30

# 将数据装载镜data中, 对数据进行分批训练
torch_data = Data.TensorDataset(X_train, y_train)
loader = Data.DataLoader(dataset=torch_data, batch_size=BATCH_SIZE, shuffle=True)


# 创建自己的nn
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        # 建立隐藏层，下面是三个隐藏层
        self.linear = nn.Linear(1,10)  # 输入是一个数，所以输出是1，中间建立10个隐藏层节点，此时，你输出的是10个数
        self.linear2 = nn.Linear(10,6,bias=True)  # 承接上一层的输入，10个数，此层在融合成6个数。这里加了个偏置参数。
        self.linear3 = nn.Linear(6,1)  # 输出层，承接上面6个数，最终输出一个数
        # 至此网络的部件定义完毕

    # 网络的部件定义好了，下面开始组装啦，设置数据如何在网络中流动，即前向传播
    def forward(self,x):  #输入的是x，就是指前面的输入层的输入
        # 设置数据流动过程，如果网络非常复杂，此处要注意输入数据的维度，经过各层之后得到的数据维度。
        x = F.relu(self.linear(x))  #首先经过一个线性层linear，然后经过一个激活函数relu
        x = F.relu(self.linear2(x))  #同上
        x = self.linear3(x)  # 最终经过最后的线性层，得到结果，结果是1维的哦
        return x  # 返回最后得到的数据


def weights_init(m):
    if isinstance(m, nn.Linear):
        init.kaiming_normal(m.weight.data)
        # init.xavier_normal(m.bias.data)


def train():
    adam_net = Net()
    # adam_net.apply(weights_init) # 对整个网络层进行参数初始化
    # 有初始化损失 xavier_normal  ：
    # {1: tensor(0.0972, grad_fn=<MseLossBackward>), 2: tensor(0.0731, grad_fn=<MseLossBackward>),
    # 3: tensor(0.0881, grad_fn=<MseLossBackward>), 4: tensor(0.1120, grad_fn=<MseLossBackward>),
    # 5: tensor(0.1012, grad_fn=<MseLossBackward>)}
    # 有初始化损失 kaiming_normal  ： 表现相对较好---
    # {1: tensor(0.1476, grad_fn=<MseLossBackward>), 2: tensor(0.0234, grad_fn=<MseLossBackward>),
    # 3: tensor(0.0162, grad_fn=<MseLossBackward>), 4: tensor(0.0170, grad_fn=<MseLossBackward>),
    # 5: tensor(0.0218, grad_fn=<MseLossBackward>)}

    # 没有初始化损失
    # {1: tensor(0.0265, grad_fn=<MseLossBackward>), 2: tensor(0.0121, grad_fn=<MseLossBackward>),
    # 3: tensor(0.0096, grad_fn=<MseLossBackward>), 4: tensor(0.0109, grad_fn=<MseLossBackward>),
    # 5: tensor(0.0104, grad_fn=<MseLossBackward>)}
    # 设置优化器和损失函数
    opt_adam = torch.optim.Adam(adam_net.parameters(), lr=LR)
    loss_func = nn.MSELoss()
    # 对数据进行分批训练
    # 在神经网络中传递完整的数据集一次是不够的，
    # 而且我们需要将完整的数据集在同样的神经网络中传递多次。
    # 但是请记住，我们使用的是有限的数据集，
    # 并且我们使用一个迭代过程即梯度下降。因此仅仅更新权重一次或者说使用一个 epoch 是不够的。
    # 比如对于一个有 2000 个训练样本的数据集。将 2000 个样本分成大小为 500 的 batch，那么完成一个 epoch 需要 4 个 iteration。
    all_loss = {}
    for epoch in range(EPOCH):
        print('epoch', epoch)
        for step, (b_x, b_y) in enumerate(loader):
            print('step', step)
            # 我们把x丢进模型，得到输出y。哇，是不是好简单，这样我们就得到结果了呢？但是不要高兴的太早，我们只是把输入数据放到一个啥都不懂（参数没有训练）的模型中，得到的结果肯定不准啊。不准的结果怎么办，看下一步。
            pre = adam_net(b_x)
            # 拿到模型输出的结果，我们就要看看模型算的准不准，就是计算损失函数了。
            loss = loss_func(pre, b_y)
            # 好了好了，我已经知道模型算的准不准了，那么就该让模型自己去朝着好的方向优化了。模型，你已经是个大孩子了，应该会自己优化的。
            ## 在优化之前，我们首先要清空优化器的梯度。因为每次循环都要靠这个优化器呢，不能翻旧账，就只算这次我们怎么优化。
            opt_adam.zero_grad()
            # 优化开始，首先，我们要把算出来的误差、损失倒着传回去。（是你们这些模块给我算的这个值，现在这个值有错误，错了这么多，返回给你们，你们自己看看自己错哪了）
            loss.backward()
            # 按照优化器的方式，一步一步优化吧。
            opt_adam.step()
            # print(loss)
            if (epoch + 1) % 100 == 0 or epoch == EPOCH-1:  # 中间每循环100次，偷偷看看结果咋样。
                print('Epoch[{}/{}],loss:{:.6f}'.format(epoch + 1, EPOCH, loss.data.item()))
                torch.save(adam_net.state_dict(), "model.pkl")  # current is model.pkl

    torch.save(adam_net.state_dict(), "model.pkl")  # current is model.pkl

if __name__ == '__main__':
    train()

